package com.caimiao.miao.library.utils;

import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.KeySpec;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;

public class AES2Util {
    private final static int MODE_ENCRYPTION = 1;
    private final static int MODE_DECRYPTION = 2;
    private final static String AES_KEY = "[;kz,s1!9o-^=08#";//AES 秘钥key，必须为16/24/32位

    private static byte[] encryption(int mode, byte[] content, String pwd, byte[] secret) {
        try {
            if (secret == null)
                secret = getSecret(pwd);

            SecretKey key = new SecretKeySpec(secret, "AES");
            Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
            byte[] iv = new byte[cipher.getBlockSize()];
            System.arraycopy(secret, 0, iv, 0, iv.length);
            IvParameterSpec ivParams = new IvParameterSpec(iv);
            cipher.init(mode == MODE_ENCRYPTION ? Cipher.ENCRYPT_MODE : Cipher.DECRYPT_MODE, key, ivParams);
            byte[] bytes = cipher.doFinal(content);
            if ((mode == MODE_ENCRYPTION)) {//加密
                byte[] ret = new byte[secret.length + bytes.length];
                System.arraycopy(secret, 0, ret, 0, secret.length);
                System.arraycopy(bytes, 0, ret, secret.length, bytes.length);
                return ret;
            } else {//解密
                return bytes;
            }
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        } catch (InvalidKeyException e) {
            e.printStackTrace();
        } catch (NoSuchPaddingException e) {
            e.printStackTrace();
        } catch (BadPaddingException e) {
            e.printStackTrace();
        } catch (IllegalBlockSizeException e) {
            e.printStackTrace();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    private static byte[] getSecret(String pwd) throws NoSuchAlgorithmException, InvalidKeySpecException {
        int iterationCount = 1000;
        int keyLength = 256;
        int saltLength = keyLength / 8;
        SecureRandom random = new SecureRandom();
        byte[] salt = new byte[saltLength];
        random.nextBytes(salt);
        KeySpec keySpec = new PBEKeySpec(pwd.toCharArray(), salt, iterationCount, keyLength);
        SecretKeyFactory keyFactory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA1");
        return keyFactory.generateSecret(keySpec).getEncoded();
    }

    private static byte[] encryptByAES(int mode, byte[] source, byte[] secret) {
        return encryption(mode, source, AES_KEY, secret);
    }

    /**
     * AES 加密
     *
     * @param source 需要加密的文件
     */
    public static byte[] encryptByAES(byte[] source) {
        return encryptByAES(MODE_ENCRYPTION, source, null);
    }

    /**
     * AES 解密
     *
     * @param source 需要解密的文件
     */
    public static byte[] decryptByAES(byte[] source, byte[] secret) {
        return encryptByAES(MODE_DECRYPTION, source, secret);
    }
}
